"""
Copyright 2023 Google LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

     https://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

"""Input pipeline using Huggingface datasets."""

import ml_collections
import jax
import datasets
import transformers
import grain.python as grain

from input_pipeline import _input_pipeline_utils
import multihost_dataloading


def get_datasets(config: ml_collections.ConfigDict):
  """Load huggingface dataset"""
  train_ds = datasets.load_dataset(
      config.hf_path,
      data_dir=config.hf_data_dir,
      data_files=config.hf_data_files,
      split="train",
      streaming=True,
      token=config.hf_access_token,
  )
  return train_ds, None


def preprocess_dataset(
    config: ml_collections.ConfigDict,
    dataloading_host_index,
    dataloading_host_count,
    global_mesh,
    dataset,
    add_bos=True,
    add_eos=True,
    packing=True,
    shift=True,
    num_threads=1,
):
  """preprocess dataset"""
  # Set global batch size.
  global_batch_size = config.global_batch_size_to_load

  assert global_batch_size % global_mesh.size == 0, "Batch size should be divisible number of global devices."

  if config.enable_data_shuffling:
    dataset = dataset.shuffle(seed=config.data_shuffle_seed)

  tokenizer = transformers.AutoTokenizer.from_pretrained(
      config.tokenizer_path,
      add_bos_token=add_bos,
      add_eos_token=add_eos,
      model_max_length=config.max_target_length,
      legacy=False,
  )

  dataset = dataset.map(
      _input_pipeline_utils.tokenization,
      batched=True,
      fn_kwargs={"tokenizer": tokenizer, "max_length": config.max_target_length - 1},
  )
  dataset = dataset.select_columns(["input_ids"])

  dataset = _input_pipeline_utils.HFDataSource(dataset, dataloading_host_index, dataloading_host_count, num_threads)

  operations = []
  operations.append(_input_pipeline_utils.HFNormalizeFeatures())

  if packing:
    operations.append(
        grain.experimental.PackAndBatchOperation(
            batch_size=global_batch_size // jax.process_count(),
            length_struct={"inputs": config.max_target_length, "targets": config.max_target_length},
        )
    )
    operations.append(_input_pipeline_utils.ReformatPacking())
  else:
    operations.append(_input_pipeline_utils.PadToMaxLength(config.max_target_length))
    operations.append(grain.Batch(batch_size=global_batch_size // jax.process_count(), drop_remainder=True))

  if shift:
    operations.append(_input_pipeline_utils.ShiftData(axis=1))

  # Since HuggingFace IterableDataset does not support access through index
  # Indexes generated by dummy_index_sampler is not used.
  # dummy_index_sampler is used as an input place holder for grain.Dataloader
  dummy_index_sampler = grain.IndexSampler(
      num_records=len(dataset),
      num_epochs=1,
      shard_options=grain.ShardOptions(
          shard_index=dataloading_host_index, shard_count=dataloading_host_count, drop_remainder=False
      ),
      shuffle=False,
      seed=0,
  )

  dataloader = grain.DataLoader(
      data_source=dataset,
      operations=operations,
      sampler=dummy_index_sampler,
      worker_count=1,  # only supports one worker for now, more workers results in duplicated data
      worker_buffer_size=1,
      read_options=grain.ReadOptions(num_threads=num_threads, prefetch_buffer_size=128),
  )

  train_iter = multihost_dataloading.MultiHostDataLoadIterator(dataloader, global_mesh)

  # Return multi-host jax.Array prep iterator
  return train_iter, None, None
